import gymnasium as gym
from gymnasium import spaces
import numpy as np

class SwimmerPIDEnv(gym.Env):
    def __init__(self, render_mode=None):
        super(SwimmerPIDEnv, self).__init__()
        # 创建 Swimmer-v5 环境
        self.env = gym.make('Swimmer-v5', render_mode=render_mode)
        self.render_mode = render_mode
        # 定义自定义环境的动作空间，6 个 PID 参数，每个参数范围为 [0, 10]
        self.action_space = spaces.Box(low=np.array([0, 0, 0, 0, 0, 0]), high=np.array([10, 10, 10, 10, 10, 10]), dtype=np.float32)
        # 定义自定义环境的观测空间，2 个状态变量，每个变量范围为 [-10, -10]
        self.observation_space = spaces.Box(low=np.array([-10, -10]), high=np.array([10, 10]), dtype=np.float32)
        # 目标值曲线
        self.setpoint_curve = lambda t: np.array([20/180*3.14*np.sin(t/10), 20/180*3.14*np.cos(t/10)])
        # 当前状态
        self.state = np.zeros(2)
        # 上一次的误差
        self.last_error = np.zeros(2)
        # 积分项
        self.integral = np.zeros(2)
        # 时间步
        self.timestep = 0
        self.control_output = np.zeros(2)  # 初始化控制输出

    def set_target_curve(self, curve_function):
        """设置目标曲线"""
        self.setpoint_curve = curve_function

    def reset(self, seed=None, options=None):
        # 重置 Swimmer 环境
        self.env.reset(seed=seed)
        # 随机初始化当前状态
        self.state = np.random.uniform(-1, 1, size=2)
        # 重置误差和积分项
        self.last_error = np.zeros(2)
        self.integral = np.zeros(2)
        # 重置时间步
        self.timestep = 0
        return self.state, {}

    def step(self, action):
        # 提取 PID 参数
        kp1, ki1, kd1, kp2, ki2, kd2 = action
        kp = np.array([kp1, kp2])
        ki = np.array([ki1, ki2])
        kd = np.array([kd1, kd2])
        
        # 更新目标值
        self.setpoint = self.setpoint_curve(self.timestep)
        
        # 计算误差、积分和导数项
        error = self.setpoint - self.state
        self.integral += error
        derivative = error - self.last_error
        self.last_error = error
        # 计算控制量
        self.control_output = kp * error + ki * self.integral + kd * derivative
        
        # 将控制量应用到 Swimmer 环境的两个关节
        swimmer_action = np.clip(self.control_output, self.env.action_space.low, self.env.action_space.high)
        obs, reward, done, truncated, info = self.env.step(swimmer_action)
        
        # 更新状态
        self.state = obs[1:3]  # 使用观测空间中的第 2 和第 3 个元素作为关节角度（关节1和关节2的角度）
        # 奖励为负的绝对误差之和
        reward = -np.sum(np.abs(error))
        # 当所有误差小于阈值时，认为任务完成
        done = np.all(np.abs(error) < 0.01)
        
        # 增加时间步
        self.timestep += 1
        
        return self.state, reward, done, truncated, info

    def get_control_output(self):
        """获取控制输出"""
        return self.control_output

    def render(self):
        # 渲染 Swimmer 环境
        return self.env.render()

    def close(self):
        # 关闭 Swimmer 环境
        self.env.close()

if __name__ == "__main__":
    env = SwimmerPIDEnv(render_mode='human')
    state, _ = env.reset()
    for _ in range(1000):
        action = np.random.uniform(low=env.action_space.low/100.0, high=env.action_space.high/100.0)  # 生成更大的随机动作
        state, reward, done, truncated, info = env.step(action)
        print(f"Action: {action}")  # 打印动作
        print(f"State: {state}")  # 打印状态
        print(f"Reward: {reward}")  # 打印奖励
        env.render()
        if done:
            state, _ = env.reset()  # 重新初始化环境状态
    env.close()
